from typing import List
from pfa import Graph, LinkInfo, solve
import random
import time


port_bandwitdth = 10000
num_gpu = 8192
num_leaf_switches = 256
num_spine_switches = 256
num_nodes = num_gpu + num_leaf_switches + num_spine_switches
num_flows = 8192


def gpu_id(i: int) -> int:
    return i


def leaf_switch_id(i: int) -> int:
    return i + num_gpu


def leaf_switch_id_by_gpu(i: int) -> int:
    return leaf_switch_id(i // num_gpu * num_leaf_switches)


def spine_switch_id(i: int) -> int:
    return i + num_gpu + num_leaf_switches


def make_graph() -> Graph:

    g = Graph(num_nodes)

    for i in range(num_gpu):
        a = gpu_id(i)
        b = leaf_switch_id_by_gpu(i)
        g.add_edge(a, b, LinkInfo(capacity=port_bandwitdth))
        g.add_edge(b, a, LinkInfo(capacity=port_bandwitdth))

    for i in range(num_leaf_switches):
        a = leaf_switch_id(i)
        for j in range(num_spine_switches):
            b = spine_switch_id(j)
            g.add_edge(a, b, LinkInfo(capacity=port_bandwitdth))
            g.add_edge(b, a, LinkInfo(capacity=port_bandwitdth))
    return g


def make_random_flow() -> List[int]:
    src = random.randint(0, num_gpu-1)
    dst = random.randint(0, num_gpu-1)
    spine = random.randint(0, num_spine_switches-1)
    src_id = gpu_id(src)
    dst_id = gpu_id(dst)
    spine_id = spine_switch_id(spine)
    leaf_id_1 = leaf_switch_id_by_gpu(src)
    leaf_id_2 = leaf_switch_id_by_gpu(dst)
    return [src_id, leaf_id_1, spine_id, leaf_id_2, dst_id]


t = 0
for i in range(1):
    g = make_graph()
    flows = {i: make_random_flow() for i in range(num_flows)}
    start = time.time()
    solution = solve(g, flows)
    t += time.time() - start
    # print(solution)
print(t)
